{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from sklearn.metrics import classification_report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "VOCAB_SIZE = 20000\n",
    "EMBED_DIM = 128\n",
    "RNN_SIZE = 128\n",
    "CLIP_NORM = 5.0\n",
    "BATCH_SIZE = 32\n",
    "DISPLAY_STEP = 50\n",
    "N_EPOCH = 2\n",
    "N_CLASS = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sort_by_len(x, y):\n",
    "    idx = sorted(range(len(x)), key=lambda i: len(x[i]))\n",
    "    return x[idx], y[idx]\n",
    "\n",
    "(X_train, y_train), (X_test, y_test) = tf.keras.datasets.imdb.load_data(num_words=VOCAB_SIZE)\n",
    "X_train, y_train = sort_by_len(X_train, y_train)\n",
    "X_test, y_test = sort_by_len(X_test, y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_sentence_batch(sent_batch):\n",
    "    max_seq_len = max([len(sent) for sent in sent_batch])\n",
    "    padded_seqs = [(sent + [0]*(max_seq_len - len(sent))) for sent in sent_batch]\n",
    "    return padded_seqs\n",
    "\n",
    "def next_train_batch():\n",
    "    for i in range(0, len(X_train), BATCH_SIZE):\n",
    "        padded_seqs = pad_sentence_batch(X_train[i : i+BATCH_SIZE])\n",
    "        yield padded_seqs, y_train[i : i+BATCH_SIZE]\n",
    "        \n",
    "def next_test_batch():\n",
    "    for i in range(0, len(X_test), BATCH_SIZE):\n",
    "        padded_seqs = pad_sentence_batch(X_test[i : i+BATCH_SIZE])\n",
    "        yield padded_seqs\n",
    "        \n",
    "def pipeline(mode):\n",
    "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        dataset = tf.data.Dataset.from_generator(next_train_batch, (tf.int32,tf.int64),\n",
    "            (tf.TensorShape([None,None]),tf.TensorShape([None])))\n",
    "        dataset = dataset.repeat(N_EPOCH)\n",
    "    if mode == tf.estimator.ModeKeys.PREDICT:\n",
    "        dataset = tf.data.Dataset.from_generator(next_test_batch, tf.int32,\n",
    "            tf.TensorShape([None,None]))\n",
    "    iterator = dataset.make_one_shot_iterator()\n",
    "    return iterator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rnn_cell():\n",
    "    return tf.nn.rnn_cell.GRUCell(RNN_SIZE//2, kernel_initializer=tf.orthogonal_initializer())\n",
    "\n",
    "def forward(inputs, reuse, is_training):\n",
    "    with tf.variable_scope('model', reuse=reuse):\n",
    "        x = tf.contrib.layers.embed_sequence(inputs, VOCAB_SIZE, EMBED_DIM)\n",
    "        x = tf.layers.dropout(x, 0.2, training=is_training)\n",
    "        _, bi_states = tf.nn.bidirectional_dynamic_rnn(\n",
    "            rnn_cell(), rnn_cell(), x, tf.count_nonzero(inputs, 1), dtype=tf.float32)\n",
    "        x = tf.concat(bi_states, -1)\n",
    "        logits = tf.layers.dense(x, N_CLASS)\n",
    "    return logits\n",
    "\n",
    "def clip_grads(loss):\n",
    "    params = tf.trainable_variables()\n",
    "    grads = tf.gradients(loss, params)\n",
    "    clipped_grads, _ = tf.clip_by_global_norm(grads, CLIP_NORM)\n",
    "    return zip(clipped_grads, params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "ops = {}\n",
    "\n",
    "X_train_batch, y_train_batch = pipeline('train').get_next()\n",
    "\n",
    "logits_train_batch = forward(X_train_batch, reuse=False, is_training=True)\n",
    "\n",
    "ops['global_step'] = tf.Variable(0, trainable=False)\n",
    "\n",
    "ops['lr'] = tf.train.exponential_decay(5e-3, ops['global_step'], 1400, 0.2)\n",
    "\n",
    "ops['loss'] = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
    "    logits=logits_train_batch, labels=y_train_batch))\n",
    "\n",
    "ops['train'] = tf.train.AdamOptimizer(ops['lr']).apply_gradients(\n",
    "    clip_grads(ops['loss']), global_step=ops['global_step'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 1 | Loss 0.714 | LR: 0.0050\n",
      "Step 50 | Loss 0.349 | LR: 0.0047\n",
      "Step 100 | Loss 0.280 | LR: 0.0045\n",
      "Step 150 | Loss 0.379 | LR: 0.0042\n",
      "Step 200 | Loss 0.304 | LR: 0.0040\n",
      "Step 250 | Loss 0.246 | LR: 0.0038\n",
      "Step 300 | Loss 0.368 | LR: 0.0035\n",
      "Step 350 | Loss 0.357 | LR: 0.0033\n",
      "Step 400 | Loss 0.654 | LR: 0.0032\n",
      "Step 450 | Loss 0.452 | LR: 0.0030\n",
      "Step 500 | Loss 0.369 | LR: 0.0028\n",
      "Step 550 | Loss 0.323 | LR: 0.0027\n",
      "Step 600 | Loss 0.153 | LR: 0.0025\n",
      "Step 650 | Loss 0.104 | LR: 0.0024\n",
      "Step 700 | Loss 0.204 | LR: 0.0022\n",
      "Step 750 | Loss 0.195 | LR: 0.0021\n",
      "Step 800 | Loss 0.071 | LR: 0.0020\n",
      "Step 850 | Loss 0.065 | LR: 0.0019\n",
      "Step 900 | Loss 0.148 | LR: 0.0018\n",
      "Step 950 | Loss 0.039 | LR: 0.0017\n",
      "Step 1000 | Loss 0.139 | LR: 0.0016\n",
      "Step 1050 | Loss 0.089 | LR: 0.0015\n",
      "Step 1100 | Loss 0.053 | LR: 0.0014\n",
      "Step 1150 | Loss 0.195 | LR: 0.0013\n",
      "Step 1200 | Loss 0.003 | LR: 0.0013\n",
      "Step 1250 | Loss 0.204 | LR: 0.0012\n",
      "Step 1300 | Loss 0.109 | LR: 0.0011\n",
      "Step 1350 | Loss 0.231 | LR: 0.0011\n",
      "Step 1400 | Loss 0.016 | LR: 0.0010\n",
      "Step 1450 | Loss 0.029 | LR: 0.0009\n",
      "Step 1500 | Loss 0.081 | LR: 0.0009\n"
     ]
    }
   ],
   "source": [
    "sess = tf.Session()\n",
    "sess.run(tf.global_variables_initializer())\n",
    "while True:\n",
    "    try:\n",
    "        sess.run(ops['train'])\n",
    "    except tf.errors.OutOfRangeError:\n",
    "        break\n",
    "    else:\n",
    "        step = sess.run(ops['global_step'])\n",
    "        if step % DISPLAY_STEP == 0 or step == 1:\n",
    "            loss, lr = sess.run([ops['loss'], ops['lr']])\n",
    "            print(\"Step %d | Loss %.3f | LR: %.4f\" % (step, loss, lr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "ops['predict'] = forward(pipeline('infer').get_next(), reuse=True, is_training=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.8980\n",
      "             precision    recall  f1-score   support\n",
      "\n",
      "          0       0.88      0.92      0.90     12500\n",
      "          1       0.92      0.88      0.90     12500\n",
      "\n",
      "avg / total       0.90      0.90      0.90     25000\n",
      "\n"
     ]
    }
   ],
   "source": [
    "y_pred_li = []\n",
    "while True:\n",
    "    try:\n",
    "        y_pred_li.append(sess.run(ops['predict']))\n",
    "    except tf.errors.OutOfRangeError:\n",
    "        break\n",
    "y_pred = np.argmax(np.vstack(y_pred_li), 1)\n",
    "print(\"Accuracy: %.4f\" % (y_pred==y_test).mean())\n",
    "print(classification_report(y_test, y_pred))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
